package org.apache.mina.protocol.dns;

import java.net.InetAddress;
import java.net.SocketAddress;
import java.net.UnknownHostException;
import java.util.ArrayList;
import java.util.List;
import java.util.Queue;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;

import org.apache.mina.common.ConnectFuture;
import org.apache.mina.common.IoSession;
import org.apache.mina.protocol.dns.messages.DnsMessage;
import org.apache.mina.protocol.dns.nameserver.ManualNameServerConfiguration;
import org.apache.mina.protocol.dns.nameserver.NameServerConfiguration;
import org.apache.mina.protocol.dns.nameserver.NameServerConfigurationFactory;
import org.apache.mina.transport.socket.DatagramConnector;
import org.apache.mina.transport.socket.nio.NioDatagramConnector;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Asynchronous DNS resolver
 */
public class DnsResolver {
    private static Logger log = LoggerFactory.getLogger(DnsResolver.class);

    private IoSession dnsSession;

    private ConcurrentMap<Integer, Queue<DnsRequest>> pendingRequests = new ConcurrentHashMap<Integer, Queue<DnsRequest>>();

    private ConcurrentMap<String, Integer> pendingTransactions = new ConcurrentHashMap<String, Integer>();

    private DnsCache dnsCache = new DnsCache();

    private Random random = new Random();

    private DnsRetry dnsRetry = new DnsRetry();

    /**
     * Default constructor.
     * 
     * Constructs a {@link DnsResolver} using default name servers and domains 
     * as configured in the host machine. 
     */
    public DnsResolver() {
        this(NameServerConfigurationFactory.getNameServers());
    }

    /**
     * Constructs a {@link DnsResolver} using a custom 
     * {@link NameServerConfiguration}.
     * 
     * Can be used to force the {@link DnsResolver} to use non-default
     * name servers:
     * 
     * <pre>
     * <code>
     * {@link NameServerConfiguration} conf = new {@link ManualNameServerConfiguration}();
     * conf.addNameServer( "10.0.0.1" );
     * conf.addDomain( "mydomain.com" );
     * 
     * {@link DnsResolver} resolver = new {@link DnsResolver}( conf );
     * </code>
     * </pre>
     * 
     * @param nameServerConfiguration {@link NameServerConfiguration} object 
     *        that holds the resolver configuration
     *        
     *  @see ManualNameServerConfiguration
     */
    public DnsResolver(NameServerConfiguration nameServerConfiguration) {
        DatagramConnector connector = new NioDatagramConnector();
        connector.setHandler(new DnsHandler(this));

        log.debug("Using DNS configuration from {}", nameServerConfiguration
                .getClass().getSimpleName());
        log.debug("Name servers: {}", nameServerConfiguration.getNameServers());
        log.debug("Search domains: {}", nameServerConfiguration.getDomains());

        SocketAddress nameServerAddress = nameServerConfiguration
                .getNameServers().get(0);
        ConnectFuture connFuture = connector.connect(nameServerAddress);

        connFuture.awaitUninterruptibly();

        dnsSession = connFuture.getSession();
        dnsRetry.start();
    }

    /**
     * Asynchronously resolves a domain name into an IP address. 
     * 
     * 
     * @param host the hostname to resolve into an IP address
     * @return a {@link DnsFuture} which can be used to be notified 
     *         when the DNS lookup is finished
     */
    public DnsFuture resolve(String host) {
        log.debug("Tring to resolve '{}'", host);

        List<InetAddress> addresses = dnsCache.get(host);
        if (addresses != null) {
            log.debug("Found in cache: {}", addresses);
            return new DnsFutureResolved(host, addresses);
        }

        if (isIpAddress(host)) {
            log.debug("Host is an IP address");
            try {
                InetAddress addr = InetAddress.getByName(host);
                addresses = new ArrayList<InetAddress>(1);
                addresses.add(addr);
                dnsCache.put(host, addresses, DnsCache.MAX_TTL);
                log.debug("Created address: {}", addr);
                return new DnsFutureResolved(host, addresses);

            } catch (UnknownHostException e) {
                DnsFutureResolved future = new DnsFutureResolved(host, null);
                future.setNameError();
                return future;
            }

        } else {
            Integer transactionId = pendingTransactions.get(host);
            if (transactionId != null) {
                // Another DNS request is pending for the same host, avoid
                // sending a new request.
                DnsRequest dnsRequest = new DnsRequest(transactionId, host);
                pendingRequests.get(transactionId).add(dnsRequest);
                return dnsRequest.getFuture();

            } else {
                // Make a new DNS request for the host
                DnsRequest dnsRequest = new DnsRequest(newTransactionId(), host);
                Queue<DnsRequest> requestList = new ConcurrentLinkedQueue<DnsRequest>();
                requestList.add(dnsRequest);
                pendingRequests.put(dnsRequest.getTransactionId(), requestList);
                dnsSession.write(DnsMessage.newDnsQuery(dnsRequest
                        .getTransactionId(), host));

                return dnsRequest.getFuture();
            }
        }
    }

    protected void receivedResponse(int transactionId,
            List<InetAddress> addresses, int ttl) {
        Queue<DnsRequest> requestList = pendingRequests.remove(transactionId);
        pendingTransactions.remove(requestList.peek().getHost());

        dnsCache.put(requestList.peek().getHost(), addresses, ttl);

        log.debug("Resolved: {} : {}", requestList.peek().getHost(), addresses);

        for (DnsRequest request : requestList) {
            request.done(addresses);
        }
    }

    protected void nameError(int transactionId) {
        Integer key = Integer.valueOf(transactionId);
        Queue<DnsRequest> requestList = pendingRequests.get(key);
        pendingRequests.remove(key);
        pendingTransactions.remove(requestList.peek().getHost());

        log.debug("Name error: {} : {}", requestList.peek().getHost());
        for (DnsRequest request : requestList) {
            request.nameError();
        }
    }

    private int newTransactionId() {
        int id;
        do {
            id = random.nextInt(0xFFFF);
        } while (pendingRequests.containsKey(Integer.valueOf(id)));

        return id;
    }

    private static final String IP_PATTERN = "^((\\d{1,2}|1\\d\\d|2[0-4]\\d|25[0-5])\\.){3}(\\d{1,2}|1\\d\\d|2[0-4]\\d|25[0-5])$";

    private boolean isIpAddress(String host) {
        return host.matches(IP_PATTERN);

    }

    private class DnsRetry extends Thread {

        @Override
        public void run() {
            while (true) {
                try {
                    sleep(500);
                } catch (Exception e) {
                    return;
                }

                long now = System.currentTimeMillis();

                for (Queue<DnsRequest> reqList : pendingRequests.values()) {
                    DnsRequest req = reqList.peek();

                    if (req.getLastRequestedTime() + 1000 > now)
                        continue;

                    if (req.getAndSetRetries() > 5) {
                        // Reached the maximum retry number
                        pendingRequests.remove(req.getTransactionId());
                        pendingTransactions.remove(req.getHost());
                        req.nameError();
                    }

                    req.setLastRequestedTime(now);
                    dnsSession.write(DnsMessage.newDnsQuery(req
                            .getTransactionId(), req.getHost()));
                }
            }
        }

    }
}
